from collections import Counter

from .myseq import MySeq
from .align import SubstMatrix,Table
from .bio_sm import SubstMatrix_Bio

class MyAlign(object):
    def __init__(self,
                 seqs:list[str],
                 align_type:str= "protein") -> None:
        '''
        align_type: str["DNA"|"RNA"|"protein"]
        '''
        self.seqs = seqs
        self.align_type = align_type

    def __len__(self):
        return len(self.seqs[0])
    
    def __getitem__(self, args):
        if isinstance(args,int):
            return self.seqs[args]
        elif isinstance(args,tuple) and len(args) == 2:
            return self.seqs[args[0]][args[1]]
        else:
            raise TypeError(f"takes more than 2 positional argument")

    def __str__(self):
        return "\n".join(self.seqs)

    def num_seqs(self):
        return len(self.seqs)

    def column(self,index):
        return [seq[index] for seq in self.seqs]
    
    def __most_frequent(lst):
        
        counter = Counter((char for char in lst if char != "-"))
        if counter:
            return counter.most_common(1)[0][0]
        else:
            raise ValueError("only gap in one position")
    
    def consensus(self) -> str:
        '''
        the most frequent char in each column, ignoring gaps
        '''
        return "".join((MyAlign.__most_frequent(self.column(i))
                        for i in range(len(self))))

class PairwiseAlignment(object):
    '''
    MySeq + SubstMatrix + Table
    '''
    def __init__(self,sm:SubstMatrix,gap:int) -> None:
        self.table = None
        self.seq1 = None
        self.seq2 = None
        self.sm = sm
        self.gap = gap

    def score_pos(self,c1:str,c2:str) -> int:
        if c1 == '-' or c2 == '-':
            return self.gap
        else:
            return self.sm[c1,c2]

    def score_alin(self,align:MyAlign) -> int:
        score = 0
        for index in range(len(align)):
            c1,c2 = align.column(index)
            score += self.score_pos(c1,c2)
        return score
    
    def __validate_seq_type(self,seq1:MySeq,seq2:MySeq) -> None:
        if seq1.seq_type != seq2.seq_type:
            raise ValueError("seq1 is {} but seq2 is {}".format(
                             seq1.seq_type,seq2.seq_type))
        self.seq1 = seq1
        self.seq2 = seq2

    def needleman_Wunsch(self,seq1:MySeq,seq2:MySeq) -> int:
        self.__validate_seq_type(seq1,seq2)
        self.table = Table(row=self.seq1.seq,colum=self.seq2.seq)
        self.table.needleman_Wunsch(self.sm,self.gap)
        return self.table[-1,-1].value

    def recover_align(self) -> MyAlign:
        aligned_seqs = self.table.recover_align()
        return MyAlign(aligned_seqs,self.seq1.seq_type)

    def swith_Watermen(self,seq1:MySeq,seq2:MySeq) -> int:
        self.__validate_seq_type(seq1,seq2)
        self.table = Table(row=self.seq1.seq,colum=self.seq2.seq)
        max_cell =  self.table.smith_Waterman(self.sm,self.gap)[-1]
        return max_cell.value

    def recover_align_local(self):
        aligned_seqs = self.table.recover_align_local()[-1]
        return MyAlign(aligned_seqs,self.seq1.seq_type)

class MultipleAlignment(object):
    def __init__( self, seqs:list[MySeq], aligner:PairwiseAlignment):
        self.seqs = seqs
        self.align_pars = aligner
    
    def add_seq_alignment(self, alignment:MyAlign, seq:MySeq):
        res = ["" for _ in range(len(alignment.seqs))]
        cons = MySeq(alignment.consensus(),alignment.align_type)
        self.align_pars.needleman_Wunsch(cons, seq)
        align2 = self.align_pars.recover_align()
        orig = 0
        for i in range(len(align2)):
            if align2[0,i]== '-':
                for k in range(len(alignment.seqs)):
                    res[k] += '-'
            else:
                for k in range(len(alignment.seqs)):
                    res[k] += alignment[k,orig]
                orig+=1
        res.append(align2.seqs[1])
        return MyAlign(res, alignment.align_type)
    
    def align_consensus(self):
        self.align_pars.needleman_Wunsch(self.seqs[0], self.seqs[1])
        alignment = self.align_pars.recover_align()
        for seq in self.seqs[2:]:
            alignment = self.add_seq_alignment(alignment,seq)
        return alignment
    
if __name__ == "__main__":
    alig = MyAlign(["ATGA-A","AA-AT-"], "dna")
    print(alig)
    print(len(alig))
    print(alig.column(2))
    print(alig[1,1])
    print(alig[0])
    print(alig.consensus())

    s1 = MySeq("ATAGC")
    s2 = MySeq("AACC")
    s3 = MySeq("ATGAC")
    sm = SubstMatrix()
    sm.create_submat(1,-1,MySeq.ALPHABET['DNA'])
    aseq = PairwiseAlignment(sm,-1)
    ma = MultipleAlignment([s1,s2,s3], aseq)
    print(ma.align_consensus()) 

    seq1 = MySeq("MHQAIFIYQIGYPLKSGYIQSIRSPEYDNW")
    seq2 = MySeq("MHIFIYQIGYALKSGYIQSIRSPEYNW")
    seq3 = MySeq("MHQAIFIQIGYALKSGYQSIRSPEYDNW")
    sm = SubstMatrix_Bio("BLOSUM62")
    aseq = PairwiseAlignment(sm,-10)
    ma = MultipleAlignment([seq1,seq2,seq3], aseq)
    print(ma.align_consensus()) 
